package com.yupi.oj.manager;

import cn.hutool.core.collection.CollUtil;
import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.model.context.chat.ContextChatCompletionRequest;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.Flowable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

/**
 * AI操作管理
 */
@Component
public class AiManager {
    @Autowired
    private ArkService arkService;
    private static final Double STABLE_TEMPERATURE = 0.2;
    private static final Double UNSTABLE_TEMPERATURE = 0.8;
    //调用随机性的非流式AI
    public String doUnStableUnStreamAi(String systemMessage,String userMessage){
        return doUnStreamAi(systemMessage,userMessage,UNSTABLE_TEMPERATURE);
    }
    //调用稳定性的非流式Ai
    public String doStableUnStreamAi(String systemMessage,String userMessage){
        return doUnStreamAi(systemMessage,userMessage,STABLE_TEMPERATURE);
    }
    //接受多个用户信息与代理信息
    public String doAnyUnStreamAi(String systemMessage,List<String> userMessage, List<String> assistantMessage, Double temperature){
        List<ChatMessage> chatMessages = new ArrayList<>();
        ChatMessage chatSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(systemMessage).build();
        chatMessages.add(chatSystemMessage);
        for(int i=0;i<userMessage.size();i++){
            if(i == userMessage.size()-1){
                ChatMessage chatUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userMessage.get(i)).build();
                chatMessages.add(chatUserMessage);
                break;
            }
            ChatMessage chatUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userMessage.get(i)).build();
            chatMessages.add(chatUserMessage);
            ChatMessage chatAssistantMessage = ChatMessage.builder().role(ChatMessageRole.ASSISTANT).content(assistantMessage.get(i)).build();
            chatMessages.add(chatAssistantMessage);
        }
        return doUnStreamAi(chatMessages,temperature);
    }
    public String doUnStreamAi(String systemMessage,String userMessage,Double temperature){
        List<ChatMessage> chatMessages = new ArrayList<>();
        ChatMessage chatSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(systemMessage).build();
        ChatMessage chatUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userMessage).build();
        chatMessages.add(chatSystemMessage);
        chatMessages.add(chatUserMessage);
        return doUnStreamAi(chatMessages,temperature);
    }
    //非流式调用
    public String doUnStreamAi(List<ChatMessage> messages,Double temperature){
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("deepseek-v3-241226")
                .messages(messages)
                .temperature(temperature)
                .stream(Boolean.FALSE)
                .build();
        ChatCompletionResult chatCompletion = arkService.createChatCompletion(chatCompletionRequest);
        Object content = chatCompletion.getChoices().get(0).getMessage().getContent();
        return content.toString();
    }

    //根据上下文流式调用
    public Flowable<ChatCompletionChunk> doAnyStreamAi(String systemMessage, List<String> userMessage, List<String> assistantMessage, Double temperature){
        List<ChatMessage> chatMessages = new ArrayList<>();
        ChatMessage chatSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(systemMessage).build();
        chatMessages.add(chatSystemMessage);
        for(int i=0;i<userMessage.size();i++){
            if(i == userMessage.size()-1){
                ChatMessage chatUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userMessage.get(i)).build();
                chatMessages.add(chatUserMessage);
                break;
            }
            ChatMessage chatUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userMessage.get(i)).build();
            chatMessages.add(chatUserMessage);
            if(CollUtil.isNotEmpty(assistantMessage)){
                ChatMessage chatAssistantMessage = ChatMessage.builder().role(ChatMessageRole.ASSISTANT).content(assistantMessage.get(i)).build();
                chatMessages.add(chatAssistantMessage);
            }
        }
        return doStreamAi(chatMessages,temperature);
    }
    //根据上下文流式调用
    public Flowable<ChatCompletionChunk> doStreamAi(List<ChatMessage> messages,Double temperature){
        ChatCompletionRequest contextChatCompletionRequest = ChatCompletionRequest.builder()
                .model("deepseek-v3-241226")
                .messages(messages)
                .temperature(temperature)
                .stream(Boolean.TRUE)
                .build();
        Flowable<ChatCompletionChunk> flowable = arkService.streamChatCompletion(contextChatCompletionRequest);
        return flowable;
    }
}
